from enum import Enum

import math


def _erase_cross(duration=None, total_duration=None, direction='C', type='in'):
    if total_duration is None:
        total_duration = duration

    def __trans_img(img, t):
        h, w = img.shape[:2]
        wv = 1.0 * w / (2 * duration)
        hv = 1.0 * h / (2 * duration)
        flag = 1
        if direction == 'C':
            p = math.ceil(wv * t)
            q = math.ceil(hv * t)
        else:
            p = math.ceil(wv * (duration - t))
            q = math.ceil(hv * (duration - t))
            flag = 1 - flag
        pic = img * (1 - flag)
        pic[:q, :p] = img[:q, :p] * flag
        pic[h - q:, :p] = img[h - q:, :p] * flag
        pic[:q, w - p:] = img[:q, w - p:] * flag
        pic[h - q:, w - p:] = img[h - q:, w - p:] * flag
        return pic

    return lambda gf, t: __trans_img(gf(t), t) if type == 'in' and t <= duration \
        else __trans_img(gf(t), (total_duration - t)) if type == 'out' and t >= total_duration - duration \
        else gf(t)


def _erase_direction(duration=None, total_duration=None, direction='E', type='in'):
    if total_duration is None:
        total_duration = duration
    """
    :param direction:
               +------+------+------+
              | ' ' | 'N'  | ' ' |
              +------+------+------+
              | 'W'  | 'C' 'OC'| ' E' |
              +------+------+------+
              | ' ' | 'S'  | ' ' |
              +------+------+------+
    :param duration:
    :return:
    """

    def __trans_img(img, t):
        h, w = img.shape[:2]
        wv = 1.0 * w / duration
        hv = 1.0 * h / duration
        p = math.ceil(wv * t)
        q = math.ceil(hv * t)
        flag = 1 if type == 'in' else 0
        pic = img * (1 - flag)
        if direction == 'E':
            pic[:, :p] = img[:, :p] * flag
        elif direction == 'W':
            pic[:, w - p:] = img[:, w - p:] * flag
        elif direction == 'S':
            pic[:q, :] = img[:q, :] * flag
        elif direction == 'N':
            pic[h - q:, :] = img[h - q:, :] * flag
        return pic

    return lambda gf, t: __trans_img(gf(t), t) if type == 'in' and t <= duration \
        else __trans_img(gf(t), t - (total_duration - duration)) if type == 'out' and t >= total_duration - duration \
        else gf(t)


class EraseInFunctions(Enum):
    cross = lambda duration, total_duration, direction: _erase_cross(duration, total_duration, direction, type='in')
    direction = lambda duration, total_duration, direction: _erase_direction(duration, total_duration, direction,
                                                                             type='in')


class EraseOutFunctions(Enum):
    cross = lambda duration, total_duration, direction: _erase_cross(duration, total_duration, direction, type="out")
    direction = lambda duration, total_duration, direction: _erase_direction(duration, total_duration, direction,
                                                                             type="out")
